#---------------------------------------------------------------
# eii_demo_app
# Designed by Aixi.Wang
#---------------------------------------------------------------
# v0.1 2021.1.31
# *initial version
#
# v0.2 2021.10.16
# *integrated with eii_light_lib
# *added 3 buttons to control vi stop/start/snapshot
#
# v0.3 2021.10.26
# *supported receiving raw data streaming
#
#----------------------
'''
{
    "ingestor": {
        "loop_video": false,
        "pipeline": "CAM1",
        "poll_interval": 0.1,
        "queue_size": 1,
        "type": "hikrobotcam"
    },
    "max_jobs": 1,
    "max_workers": 1,
    "sw_trigger": {
        "init_state": "running"
    },
    "udfs": [
        {
            "name": "mydummy.main",
            "type": "python"
        }
    ]
}

or 
{
    "ingestor": {
        "loop_video": false,
        "pipeline": "/dev/video0",
        "poll_interval": 0.1,
        "queue_size": 1,
        "type": "opencv"
    },
    "max_jobs": 1,
    "max_workers": 1,
    "sw_trigger": {
        "init_state": "running"
    },
    "udfs": [
        {
            "name": "mydummy.main",
            "type": "python"
        }
    ]
}


'''
# import pygame module in this program 
import pygame
from io import BytesIO
#from PIL import Image
import threading
#import eii
import sys
import time
import zmq
import os
import socket
import queue
import json
import contextlib
import numpy as np
import cv2
#sys.stdout = open(os.devnull, "w")
#sys.stderr = open(os.devnull, "w")
#os.environ["SDL_VIDEODRIVER"] = "dummy"


from eii_light_lib import *
#from VideoCapture import Device
exit_flag = 0
Q=queue.Queue(100)
Q2=queue.Queue(100)

config = load_config()

if 'zmq_addr1' in config:
    zmq_addr1 = config['zmq_addr1']
else:
    # vi service addr
    zmq_addr1 = 'tcp://192.168.2.110:64013'

if 'zmq_addr2' in config:
    zmq_addr2 = config['zmq_addr2']
else:
    # va publisher addr
    zmq_addr2 = 'tcp://192.168.2.110:65013'


print('zmq_addr1:',zmq_addr1)
print('zmq_addr2:',zmq_addr2)

IMG_W = 1024
IMG_H = 768
real_img_w = 1024
real_img_h = 768

grey_palette = []
for i in range(0,256):
    grey_palette.append([i,i,i])


#----------------------
# draw_button
#----------------------
def draw_button(screen,t,pos,mode):
    smallfont = pygame.font.SysFont('Corbel',35)

    color = (255,255,255)
    color_light = (170,170,170)
    color_dark = (100,100,100)

    text = smallfont.render(t,True,color)

    # unclicked
    if mode == 0:
        pygame.draw.rect(screen,color_light,[pos[0],pos[1],200,40])
    
    # flying on 
    if mode == 1:
        pygame.draw.rect(screen,color_dark,[pos[0],pos[1],200,40])

    # clicked
    if mode == 2:
        pygame.draw.rect(screen,color_light,[pos[0],pos[1],200,40])
        
    screen.blit(text,(pos[0],pos[1]))
    #pygame.display.update()

#----------------------
# mouse_in_range
#----------------------
def mouse_in_range(pos,rect):
    #print('pos:',pos)
    #print('rect:',rect)
    if pos[0] > rect[0] and pos[0] < (rect[0] + rect[2]) and pos[1] > rect[1] and pos[1] < (rect[1] + rect[3]):
        #print('in range ...')
        return 1
    else:
        return 0

    
#----------------------
# get_vm_ip
#----------------------
def get_vm_ip(port):
    while True:
        print('try to get vm ip ...')
        # get target ip from vm
        outs = os.popen('multipass list').read()
        outs = outs.split('\n')[1:]
        for d in outs:
            d2 = d.split(' ')
            for d3 in d2:
                if len(d3) >=10 and len(d3) <= 15:
                    try:
                        print('try to test ip:',d3)
                        HOST,PORT= d3,port
                        s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
                        s.settimeout(3)
                        print('try to connect:',HOST,PORT)
                        s.connect((HOST,PORT))
                        return d3
                    except Exception as e:
                        print('get_vm_ip:',str(e))
                        pass
        time.sleep(3) 

#----------------------
# display_text
#----------------------
def display_text(t,center_x,center_y,disp):
    BLACK = (0,0,0)
    GREEN = (0,255,0)
    BLUE = (0,0,128)
     
    fontObj = pygame.font.Font('freesansbold.ttf',12)
    textSurfaceObj = fontObj.render(t,True,GREEN,BLUE)
    textRectObj = textSurfaceObj.get_rect()
    textRectObj.center=(center_x,center_y)
    disp.blit(textSurfaceObj,textRectObj)

#----------------------
# display_text2
#----------------------
def display_text2(fontObj,t,center_x,center_y,disp):
    BLACK = (0,0,0)
    GREEN = (0,255,0)
    BLUE = (0,0,128)
     
    #fontObj = pygame.font.Font('freesansbold.ttf',12)
    textSurfaceObj = fontObj.render(t,True,GREEN,BLUE)
    textRectObj = textSurfaceObj.get_rect()
    textRectObj.center=(center_x,center_y)
    disp.blit(textSurfaceObj,textRectObj)


                            
#----------------------
# eii_recv_thread
#----------------------
def eii_recv_thread(zmq_addr,topic,pygame_surface):
    global exit_flag
    global Q2
    global IMG_H
    global IMG_W
    global real_img_h
    global real_img_w
    global grey_palette
    
    n = 0
    last_time = time.time()
    fps = 0
    
    #connect_to = sys.argv[1]
    #target_port = port
    while True:
        try:
            if exit_flag == 1:
                return
        
            #ip = get_vm_ip(target_port)
            #print('found active ip:',ip)
            #connect_to = 'tcp://' + ip + ':' + str(target_port)
            connect_to = zmq_addr
            print('connect str:',connect_to)
            
            topics = sys.argv[2:]

            ctx = zmq.Context()
            s = ctx.socket(zmq.SUB)
            ret = s.connect(connect_to)

            s.setsockopt(zmq.SUBSCRIBE,topic.encode('utf-8'))

            try:
            #if 1:
                while True:
                    if exit_flag == 1:
                        return
                        
                    msgs = s.recv_multipart()
                    #print('-----------------------------msg num:',len(msgs))           
                    #for msg in msgs:
                    #    print('recv:',msg)
                    if len(msgs) == 5:

                        n += 1
                        if n >= 10:
                            n = 0
                            t2 = time.time()
                            t1 = last_time
                            fps = 10.0/(t2-t1)
                            print('fps:',fps)
                                
                            last_time = t2
                            print('new image received:',msgs[3])
                            
                        meta_json = json.loads(msgs[3].decode('utf-8'))
                        image_bytes = msgs[4]
                        image_size = (meta_json['width'],meta_json['height'])
                        

                        print('meta_json:',meta_json)
                        if 'encoding_type' not in meta_json:
                            if meta_json['channels'] == 1:
                                #print('channels:',meta_json['channels'])
                                img_mono = np.frombuffer(image_bytes,dtype=np.uint8)
                                img_color = cv2.merge([img_mono,img_mono,img_mono])
                                
                                image = pygame.image.frombuffer(img_color, image_size,'RGB')
                                #image.set_palette(grey_palette)
                            elif meta_json['channels'] == 3:
                                #print('channels:',meta_json['channels'])
                                image = pygame.image.frombuffer(image_bytes, image_size,'BGR')
                            else:
                                print('invalid channels')
                        else:
                            print('encoded image ...')
                            image = pygame.image.load(BytesIO(image_bytes))
                        
                        
                        #writefile('t.bin',msgs[4])
                        
                        
                        new_rect = (image.get_rect().width,image.get_rect().height)
                        #print('new_rect:',new_rect)
                        real_img_w = new_rect[0]
                        real_img_h = new_rect[1]

                        q_data = [new_rect,msgs[3].decode('utf-8'),image,fps]
                        #print('emulated injector:',q_data)
                        if Q2.qsize() >= 30:
                            print('skip frame ...')
                            Q2.get()
                
                        Q2.put(q_data)
        
                        #pygame.display.set_mode(new_rect)
                        #pygame_surface.blit(image, (0, 0))
                        #display_text(msgs[3].decode('utf-8'),512,new_rect[1]+12,pygame_surface)
                        #pygame.display.update()
                    
                    #time.sleep(0.01)
            except Exception as e:
                print('exception:',str(e))
                time.sleep(1)
                
            print("Done.")
            
        except Exception as e:
            print('exception:',str(e))
            pass




#----------------------
# main
#----------------------
def main():
    global IMG_H
    global IMG_W
    global real_img_h
    global real_img_w
    
    global exit_flag
    global Q2

    pygame.init() 
    white = (255, 255, 255) 
    X = IMG_W + 250
    Y = IMG_H
    
    fontObj = pygame.font.Font('freesansbold.ttf',12)
    display_surface = pygame.display.set_mode((X,Y)) 
    
    
    width = display_surface.get_width()
    height = display_surface.get_height()

    #-----------------------------------
    # button init
    #-----------------------------------   
    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
    draw_button(display_surface,'stop ingestion ',[IMG_W + 20, 150],0)
    draw_button(display_surface,'   snapshot    ',[IMG_W + 20, 250],0)
    
    pygame.display.update() 
    
    button1_rect = [IMG_W + 20,50,200,40]
    button2_rect = [IMG_W + 20,150,200,40]
    button3_rect = [IMG_W + 20,250,200,40]

    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
    draw_button(display_surface,'stop ingestion ',[IMG_W + 20, 150],0)
    draw_button(display_surface,'   snapshot    ',[IMG_W + 20, 250],0)
    
    pygame.display.update() 
    
    button1_status = 0
    button2_status = 0
    button3_status = 0

    #-----------------------------------
    # start receiver thread
    #-----------------------------------   
    
    t1 = threading.Thread(target=eii_recv_thread,args=(zmq_addr2,'camera1_stream_results',display_surface,))
    t1.start()
    
    # set the pygame window name 
    pygame.display.set_caption('EII Demo App') 
    # create a surface object, image is drawn on it. 
    #with open('a.jpg', 'rb') as imgfile:
    #    image = pygame.image.load(BytesIO(imgfile.read()))
        
    #image = pygame.image.load(r'a.jpg') 
    #display_surface.blit(image, (0, 0))
    
    #-----------------------------------
    # main loop
    #-----------------------------------   

    #display_surface.fill(white)      
    while True : 
        # completely fill the surface object 
        # with white colour 
        if (IMG_H != real_img_h) or (IMG_W != real_img_w):
            display_surface = pygame.display.set_mode((real_img_w + 250,real_img_h))
            IMG_H = real_img_h
            IMG_W = real_img_w
    
            button1_rect = [IMG_W + 20,50,200,40]
            button2_rect = [IMG_W + 20,150,200,40]
            button3_rect = [IMG_W + 20,250,200,40]
    
            draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
            draw_button(display_surface,'stop ingestion ',[IMG_W + 20, 150],0)
            draw_button(display_surface,'   snapshot    ',[IMG_W + 20, 250],0)
            
            pygame.display.update() 
            
        for event in pygame.event.get(): 
            # if event object type is QUIT 
            # then quitting the pygame 
            # and program both. 
            if event.type == pygame.QUIT : 
                exit_flag = 1

                # deactivates the pygame library 
                pygame.quit() 
      
                # quit the program. 
                quit() 
            
            
            #-----------------------------------
            # process button event
            #-----------------------------------
            
            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = pygame.mouse.get_pos()
                #print('mouse:',pos)
                
                if mouse_in_range(pos, button1_rect):
                    resp = service_call(zmq_addr1,'default','START_INGESTION')
                    print('resp:',resp)
                    display_text2(fontObj,resp[1],IMG_W + 125,350,display_surface)
                        
                if mouse_in_range(pos, button2_rect):
                    resp = service_call(zmq_addr1,'default','STOP_INGESTION')
                    print('resp:',resp)
                    display_text2(fontObj,resp[1],IMG_W + 125,350,display_surface)
                    
                if mouse_in_range(pos, button3_rect):
                    resp = service_call(zmq_addr1,'default','SNAPSHOT')
                    print('resp:',resp)
                    display_text2(fontObj,resp[1],IMG_W + 125,350,display_surface)
                pygame.display.update() 
                
                
            #-----------------------------------
            # process button moving effect
            #-----------------------------------
            if event.type == pygame.MOUSEMOTION:
                pos = pygame.mouse.get_pos()
                
                if mouse_in_range(pos, button1_rect) == 1:
                    #print('button1 in range ...')
                    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],1)
                    draw_button(display_surface,'stop  ingestion',[IMG_W + 20, 150],0)
                    draw_button(display_surface,'snopshort      ',[IMG_W + 20, 250],0)
                    
                    button1_status = 1
                    button2_status = 0
                    button3_status = 0
                    
                if mouse_in_range(pos, button2_rect) == 1:
                    #print('button2 in range ...')                
                    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
                    draw_button(display_surface,'stop  ingestion',[IMG_W + 20, 150],1)
                    draw_button(display_surface,'snopshort      ',[IMG_W + 20, 250],0)

                    button1_status = 0
                    button2_status = 1
                    button3_status = 0


                if mouse_in_range(pos, button3_rect) == 1:
                    #print('button3 in range ...')
                    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
                    draw_button(display_surface,'stop  ingestion',[IMG_W + 20, 150],0)
                    draw_button(display_surface,'snopshort      ',[IMG_W + 20, 250],1)

                    button1_status = 0
                    button2_status = 0
                    button3_status = 1

                if (button1_status + button2_status + button3_status) > 0 \
                                and mouse_in_range(pos, button1_rect) == 0 \
                                and mouse_in_range(pos, button2_rect) == 0 \
                                and mouse_in_range(pos, button3_rect) == 0:
                    draw_button(display_surface,'start ingestion',[IMG_W + 20, 50],0)
                    draw_button(display_surface,'stop  ingestion',[IMG_W + 20, 150],0)
                    draw_button(display_surface,'snopshort      ',[IMG_W + 20, 250],0)
                    
                pygame.display.update() 
                #print('mouse:',mouse)

            #print('event:',event)

        #--------------------------------------
        # display image to screen from fifo
        #--------------------------------------
        if Q2.qsize() > 0:
            q2_data = Q2.get()
            #print('main loop, q2_data:',q2_data)
            
            if len(q2_data) == 4:

                new_rect = q2_data[0]
                new_str = q2_data[1]
                new_image = q2_data[2]
                fps = q2_data[3]
                #pygame.display.set_mode(new_rect)
                display_surface.blit(new_image, (0, 0))
                display_text2(fontObj,new_str,real_img_w/2,12,display_surface)
                display_text2(fontObj,'current fps:' + str(round(fps,1)),real_img_w/2,12 + 16,display_surface)
                #pygame.display.update()

                # Draws the surface object to the screen.   
                pygame.display.update() 


        if exit_flag == 1:
            sys.exit(0)

        
        
if __name__ == "__main__":
    main()
    #pass        